from flask import Flask, request, render_template, jsonify
from dotenv import load_dotenv
import os
from datetime import datetime
import time
from utils.qwen_client import analyze_with_qwen
from pathlib import Path
import cv2

# 加载 .env 文件
def load_environment():
    """明确加载环境变量"""
    # 确定 .env 文件路径
    env_path = Path(__file__).parent / '.env'
    print(f"Looking for .env at: {env_path}")
    
    if env_path.exists():
        print(".env file found, loading...")
        load_dotenv(dotenv_path=env_path, override=True)
        
        # 打印所有相关的环境变量
        env_vars = ['API_BASE_URL', 'API_TOKEN', 'MINIO_ACCESS_KEY',  'OLLAMA_MODEL']
        
        print("Environment variables from .env:")
        for var in env_vars:
            value = os.getenv(var)
            if value:
                print(f"  {var}: {value}")
            else:
                print(f"  {var}: NOT SET")
                
    else:
        print(".env file not found!")

# 在Hydra之前加载环境变量
load_environment()

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = '/data/tmp/'
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024  # 限制上传大小为100MB

# 允许上传的图片格式
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp4'}

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 模拟大模型识别函数（实际项目中替换为真实模型调用）
def call_ai_model(image_path):
    # 这里模拟处理时间
    # time.sleep(2)
    MODEL_TYPE = os.getenv("OLLAMA_TYPE")
    MODEL_HOST = os.getenv("OLLAMA_HOST")
    MODEL = os.getenv("OLLAMA_MODEL")
    
    result = analyze_with_qwen(MODEL_TYPE,[image_path,image_path],MODEL_HOST,MODEL,0,'')
    print(f"qwen result: {result}")
    return result

def extract_jpeg(mp4_filepath:str,jpeg_filepath) ->str:
    """提取上传的JPEG图片"""
    try:
        print(f"从 {mp4_filepath} 抽帧: {jpeg_filepath}")
        # 使用OpenCV读取视频并抽帧
        cap = cv2.VideoCapture(mp4_filepath)
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_to_capture = int(3 * fps)  # 第5秒的帧
        
        # 设置读取位置
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_to_capture)
        success, frame = cap.read()
        
        if success:
            # 保存为JPEG
            cv2.imwrite(jpeg_filepath, frame)
            print(f"从 {mp4_filepath} 成功抽帧: {jpeg_filepath}")
        else:
            # 如果视频长度不足5秒，取第一帧
            cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            success, frame = cap.read()
            if success:
                cv2.imwrite(jpeg_filepath, frame)
                print(f"视频不足5秒，取第一帧: {jpeg_filepath}")
            else:
                return jsonify({'error': '无法从视频中读取帧'})
        
        cap.release()
        
        # 删除临时MP4文件（可选）
        # os.remove(mp4_filepath)
        
    except Exception as e:
        return jsonify({'error': f'视频处理失败: {str(e)}'})
    

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': '没有选择文件'})
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'})
    
    if file and allowed_file(file.filename):
        
        filename = None
        filepath = None
        original_filename = os.path.splitext(file.filename)[0]
        if file.filename.lower().endswith('.mp4'):
            # 保存原始MP4文件
            mp4_filename = f"{original_filename}.mp4"
            mp4_filepath = os.path.join(app.config['UPLOAD_FOLDER'], mp4_filename)
            file.save(mp4_filepath)
            
            # 从MP4抽帧处理
            jpeg_filename = f"{original_filename}.jpeg"
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], jpeg_filename)
            
            extract_jpeg(mp4_filepath, filepath)
            
        else:
            # 生成时间戳文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"temp_{timestamp}.jpeg"
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            # 保存文件
            file.save(filepath)
        
        
        # 调用大模型（模拟）
        results = call_ai_model(filepath)
        
        return jsonify({
            'success': True,
            'filename': filename,
            'detection_result': results
        })
    
    return jsonify({'error': '不支持的文件格式'})

if __name__ == '__main__':
    app.run(debug=True,host='0.0.0.0',port=5000)
    # extract_jpeg("/data/tmp/test_video.mp4","/data/tmp/test_video.jpg")